#!/usr/bin/env python3
# -*- encoding: utf-8 -*-
import argparse
from uvicorn.main import Server
from uvicorn.config import Config as WebConfig
from llmweb import BrowserPool, Context, KimLLM, YiYanLLM, LoginMixin
from llmweb.app import create_app, logger
from llmweb.config import Config
from dotenv import load_dotenv
import asyncio
import logging

async def login(args, argv):
    # 从args中获取target参数
    target = args.target
    pool = BrowserPool(headless=False, workdir=Config.WORK_DIR, pool_size=1)
    # 初始化BrowserPool
    await pool.initialize()
    model: LoginMixin = None
    # 根据target选择相应的模型
    if target == 'yiyan':
        model = YiYanLLM(pool=pool)
    elif target == 'kimi':
        model = KimLLM(pool=pool)
    else:
        raise ValueError(f"Unknown target model: {target}")

    # 这里添加登录逻辑，例如调用模型的login方法
    # 假设模型有一个async的login方法，这里需要根据实际情况来编写
    await model.login()  # 假设的登录方法调用

    # 登录成功后的逻辑，例如打印信息
    print(f"Successfully logged in using model: {target}")

# async函数，用于定义异步的函数
async def serve(args, argv):
    # 创建BrowserPool实例，用于管理Browser，参数分别为headless和workdir
    pool = BrowserPool(headless=Config.HEADLESS, workdir=Config.WORK_DIR, pool_size=Config.POOL_SIZE)
    # 初始化BrowserPool
    await pool.initialize()
    # 创建一个空字典，用于存放模型
    models = dict()
    # 向字典中添加两个模型：YiYanLLM和KimLLM
    models["yiyan"] = YiYanLLM(pool=pool)
    models["kimi"] = KimLLM(pool=pool)
    # 创建Context实例，参数分别为pool和models
    ctx = Context(pool=pool, models=models)
    # 创建app，参数分别为ctx
    app = create_app(ctx)
    # 打印启动服务器的信息
    print(f"> Start API server... 0.0.0.0:{Config.PORT}")
    # 创建WebConfig实例，参数分别为app、host、port和loop
    conf = WebConfig(app=app, host="0.0.0.0", port=Config.PORT, loop='asyncio')  
    # 创建Server实例，参数分别为conf
    server = Server(config=conf)
    # 启动服务器
    await server.serve()  

async def main():
    load_dotenv()
    parser = argparse.ArgumentParser(
        description='llmweb tools')

    sub_parsers = parser.add_subparsers(help='sub commands')
    parser_serve = sub_parsers.add_parser('serve', help="start api server")
    parser_serve.set_defaults(func=serve)

    parser_init = sub_parsers.add_parser('login', help="login to web site")
    parser_init.add_argument('--target', type=str, choices=['yiyan', 'kimi'], help="Specify the target model for login", required=True)
    parser_init.set_defaults(func=login)
    
    args,argv = parser.parse_known_args()
    #setup logging
    console_handler = logging.StreamHandler()
    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    console_handler.setFormatter(formatter)
    logger.addHandler(console_handler)
    logger.setLevel(logging.DEBUG)
    logger.debug(f"begin start llmweb")
    if hasattr(args, 'func'):  
        # 使用 create_task 创建任务，而不是直接运行  
        task = asyncio.create_task(args.func(args, argv))  
        # 等待任务完成  
        await task  
    else:  
        parser.print_help()  
        return 1

if __name__ == "__main__":
    asyncio.run(main())